import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, TensorDataset, random_split, DataLoader
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import os

matplotlib.use('TkAgg')

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
non_blocking = torch.cuda.is_available()


class TitanicDataset(Dataset):
    def __init__(self, filepath):
        XY_data = np.loadtxt(filepath, delimiter=',', dtype=np.float32, skiprows=1)
        self.len = XY_data.shape[0]
        self.X = torch.from_numpy(XY_data[:, 2:])
        self.y = torch.from_numpy(XY_data[:, [1]])

    def __getitem__(self, index):
        return self.X[index], self.y[index]

    def __len__(self):
        return self.len


titanicDataset = TitanicDataset("../data/Titanic/titanic_preprocessed3.csv")
dataset = TensorDataset(titanicDataset.X, titanicDataset.y)
train_size = int(0.8 * len(dataset))
val_size = len(dataset) - train_size
train_dataset, val_dataset = random_split(dataset, [train_size, val_size])

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, pin_memory=True)


class TitanicModel(nn.Module):
    def __init__(self):
        super(TitanicModel, self).__init__()
        self.linear1 = nn.Linear(26, 24)
        self.linear2 = nn.Linear(24, 16)
        self.linear3 = nn.Linear(16, 12)
        self.linear4 = nn.Linear(12, 8)
        self.linear5 = nn.Linear(8, 6)
        self.linear6 = nn.Linear(6, 4)
        self.linear7 = nn.Linear(4, 2)
        self.linear8 = nn.Linear(2, 1)
        self.act = nn.Sigmoid()

    def forward(self, x):
        x = self.act(self.linear1(x))
        x = self.act(self.linear2(x))
        x = self.act(self.linear3(x))
        x = self.act(self.linear4(x))
        x = self.act(self.linear5(x))
        x = self.act(self.linear6(x))
        x = self.act(self.linear7(x))
        x = self.linear8(x)
        return x


titanicModel = TitanicModel().to(device)
# Binary Cross Entropy 是一种用于二分类问题的损失函数，衡量的是模型预测的概率分布与真实标签之间的差异。
# 在这个上下文中，BCEWithLogitsLoss 结合了 Sigmoid 激活函数和二元交叉熵损失函数，通常用于处理二分类任务。
criterion = nn.BCEWithLogitsLoss()
optimizer = optim.Adam(titanicModel.parameters(), lr=0.01)

epochs = 300
train_loss = []
train_acc = []
test_loss = []
test_acc = []
for epoch in range(epochs):
    # 训练阶段
    titanicModel.train()
    train_loss_epoch = 0
    train_acc_epoch = 0
    for i, (inputs, labels) in enumerate(train_loader):
        inputs = inputs.to(device, non_blocking=non_blocking)
        labels = labels.to(device, non_blocking=non_blocking)
        logits = titanicModel(inputs)
        loss = criterion(logits, labels)
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        train_loss_epoch += loss.item()
        pred = torch.sigmoid(logits)
        train_acc_epoch += ((pred > 0.5) == labels).float().mean().item()
    train_loss_epoch = train_loss_epoch / len(train_loader)
    train_acc_epoch = train_acc_epoch / len(train_loader)
    train_loss.append(train_loss_epoch)
    train_acc.append(train_acc_epoch)
    print("epoch: {}, train_loss: {}, train_acc: {}".format(epoch, train_loss_epoch, train_acc_epoch))
    # 验证阶段
    titanicModel.eval()
    test_loss_epoch = 0
    test_acc_epoch = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs = inputs.to(device, non_blocking=non_blocking)
            labels = labels.to(device, non_blocking=non_blocking)
            logits = titanicModel(inputs)
            loss = criterion(logits, labels)
            test_loss_epoch += loss.item()
            pred = torch.sigmoid(logits)
            test_acc_epoch += ((pred > 0.5) == labels).float().mean().item()
        test_loss_epoch = test_loss_epoch / len(val_loader)
        test_acc_epoch = test_acc_epoch / len(val_loader)
        test_loss.append(test_loss_epoch)
        test_acc.append(test_acc_epoch)
        print("epoch: {}, test_loss: {}, test_acc: {}".format(epoch, test_loss_epoch, test_acc_epoch))

plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
plt.plot(range(epochs), train_loss, label="train_loss")
plt.plot(range(epochs), test_loss, label="test_loss")
plt.legend()
plt.subplot(1, 2, 2)
plt.plot(range(epochs), train_acc, label="train_acc")
plt.plot(range(epochs), test_acc, label="test_acc")
plt.legend()
plt.show()

model_path = "../models"
if not os.path.exists(model_path):
    os.makedirs(model_path)
torch.save(titanicModel.state_dict(), os.path.join(model_path, "titanic_model.pth"))
